package org.xiaobai.ai.extension;

import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.spring.SpringUtil;
import lombok.extern.slf4j.Slf4j;
import org.jsoup.Connection;
import org.jsoup.Jsoup;
import org.jsoup.nodes.Document;
import org.jsoup.nodes.Element;
import org.jsoup.select.Elements;
import org.springframework.util.CollectionUtils;
import org.xiaobai.ai.api.enums.ApiKeyModelEnum;
import org.xiaobai.ai.api.enums.ApiTypeEnum;
import org.xiaobai.ai.entity.ChatMessageVO;
import org.xiaobai.core.config.ChatConfig;
import org.xiaobai.core.utils.JsoupUtil;

import java.io.IOException;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.IdentityHashMap;
import java.util.Map;

/**
 * @author dingfeng.xiao
 * @description
 * @date 2023/9/11 10:10
 */
@Slf4j
public class BingSearch {
    public static IdentityHashMap<String, String> bingSearch(String query) {
        IdentityHashMap<String,String> map = new IdentityHashMap<>();
        try {
            String bingUrl = "https://cn.bing.com/search?q=" + URLEncoder.encode(query, StandardCharsets.UTF_8) + "&qs=n&form=QBRE&sp=-1&lq=0&pq=%E4%B8%AD%E5%8D%8E%E7%83%9F%E7%8E%B0%E5%9C%A8%E5%A4%9A%E5%B0%91%E9%92%B1%E4%B8%80%E5%8C%85&sc=6-10&sk=&cvid=1079A29E914E428483AC9682F9B934A3&ghsh=0&ghacc=0&ghpl=";
            // 设置连接参数，允许重定向
            Connection.Response response = JsoupUtil.connect(bingUrl)
                    .followRedirects(true)
                    .execute();
            // 获取重定向后的页面内容
            Document document = response.parse();
            Elements links = document.select("li.b_algo a[href^=http]");
            for (Element link : links) {
                String title = link.select("h2").text();
                String linkUrl = link.attr("href");
                map.put(title, linkUrl);
            }
            log.info("搜索词：" + query);
            log.info("搜索结果：" + map.size());
        } catch (IOException e) {
            log.error("search,error:",e);
            e.printStackTrace();
        }
        return map;
    }

    public static String scrapeText(String url) {
        StringBuilder text = new StringBuilder();
        try {
            Document doc = JsoupUtil.connect(url)
                    .get();

            doc.select("script, style").remove();
            String pageText = doc.text();

            String[] lines = pageText.split("\n");
            for (String line : lines) {
                String[] chunks = line.split("  ");
                for (String chunk : chunks) {
                    chunk = chunk.trim();
                    if (!chunk.isEmpty()) {
                        text.append(chunk).append("\n");
                    }
                }
            }
        } catch (IOException e) {
            return "无法连接到该网页";
        }
        return text.toString();
    }

    public static String getChatMsg(ChatMessageVO chatMessageVO) {
        ApiTypeEnum apiTypeEnum = SpringUtil.getBean(ChatConfig.class).getApiTypeEnum();
        int maxToken = (apiTypeEnum == ApiTypeEnum.ACCESS_TOKEN ? 5500 : ApiKeyModelEnum.maxTokens(chatMessageVO.getModelName()));
        int i = 0;
        StringBuilder chatMsg = new StringBuilder("材料内容：\n");
        StringBuilder tempMsg = new StringBuilder("材料内容：\n");
        IdentityHashMap<String, String> bingSearch = bingSearch(chatMessageVO.getContent());
        if (CollectionUtils.isEmpty(bingSearch)) {
            try {
                Thread.sleep(100L);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
            bingSearch = bingSearch(chatMessageVO.getContent());
        }
        for (Map.Entry<String, String> entry : bingSearch.entrySet()) {
            if (chatMsg.length() > maxToken) {
                continue;
            }
            String scrapeText = scrapeText(entry.getValue());
            if (StrUtil.isEmpty(scrapeText)) {
                continue;
            }
            tempMsg.append(i).append("、").append(scrapeText).append("\n");
            if (tempMsg.length() > maxToken) {
                int splitIndex = maxToken - chatMsg.length();
                chatMsg.append(i).append("、").append(scrapeText, 0, Math.max(splitIndex, 0)).append("\n");
                continue;
            }
            chatMsg.append(i).append("、").append(scrapeText).append("\n");
        }
        return chatMsg.toString();
    }

}
